/*
 * Copyright 2025 Snowflake Inc.
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*
 * Files that are generated by in-progress transactions
 * are handled by the functions declared here. If a
 * transaction is successfully committed, its in-progress
 * tables are removed by the same transaction. If not
 * committed, the files are removed from remote storage
 * via VACUUM.
 */
#include "postgres.h"
#include "funcapi.h"
#include "miscadmin.h"

#include "access/table.h"
#include "access/heapam.h"
#include "access/genam.h"
#include "catalog/namespace.h"
#include "commands/sequence.h"
#include "pg_lake/cleanup/in_progress_files.h"
#include "pg_lake/extensions/pg_lake_table.h"
#include "pg_lake/extensions/pg_lake_engine.h"
#include "pg_lake/locks/advisory_locks.h"
#include "pg_lake/pgduck/remote_storage.h"
#include "pg_lake/util/array_utils.h"
#include "pg_lake/util/spi_helpers.h"
#include "pg_lake/util/plan_cache.h"
#include "pg_lake/util/string_utils.h"
#include "datatype/timestamp.h"
#include "storage/procarray.h"
#include "utils/fmgroids.h"
#include "utils/memutils.h"
#include "utils/lsyscache.h"
#include "utils/snapmgr.h"

#define OPERATION_ID_SEQUENCE "operationid_seq"


#define Natts_in_progress_files 3
#define Anum_in_progress_path 1
#define Anum_in_progress_operation_id 2
#define Anum_in_progress_is_prefix 3


/*
* Each transaction that needs unique operationId, as
* we acquire a lock on the operationId. Hence,
* we call it TransactionOperationId.
*
* The lock on the TransactionOperationId is used
* to prevent concurrent cleanup to access the same
* files generated by the transaction.
*/
static int64 TransactionOperationId = 0;

/* each backend should register the callbacks */
static bool TransactionCallbackRegistered = false;


/*
 * InProgressFiles represents a deletion entry from the
 * IN_PROGRESS_FILES_TABLE.
 */
typedef struct InProgressFiles
{
	char	   *path;
	int64_t		operationId;
	bool		isPrefix;
}			InProgressFiles;

static bool GetInProgressFileRecords(char *location, bool isFull, List **fileRecords);
static int64 GetTransactionOperationId(void);
static void InProgressFileTransactionCallback(XactEvent event, void *arg);
static int64 GetNextTransactionOperationId(void);
static bool TryLockTransactionOperationId(int64 operationId);
static void LockTransactionOperationId(int64 operationId);
static bool InProgressTableVisibleToExternalTx(void);
static bool FileInProgressRecordExists(char *path);


PG_FUNCTION_INFO_V1(flush_in_progress_queue);


/*
* flush_in_progress_queue removes all eligible files from
* the IN_PROGRESS_FILES_TABLE.
*
* Normally when vacuum kicks in for iceberg tables, it
* removes the files generated by in-progress transactions.
* User can force the cleanup by calling this function.
* The function is also useful for cleaning up the files
* that are no not iceberg tables. For example, COPY TO S3
* files that are not committed.
*/
Datum
flush_in_progress_queue(PG_FUNCTION_ARGS)
{
	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;

	InitMaterializedSRF(fcinfo, MAT_SRF_USE_EXPECTED_DESC);

	List	   *removedFiles = NIL;

	/* remove everything */
	char	   *location = "";
	bool		isVerbose = false;

	RemoveInProgressFiles(location, true, isVerbose, &removedFiles);

	ListCell   *fileCell;

	foreach(fileCell, removedFiles)
	{
		char	   *path = lfirst(fileCell);

		Datum		values[] = {CStringGetTextDatum(path)};
		bool		nulls[] = {false};

		tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc, values, nulls);
	}

	PG_RETURN_VOID();
}

/*
 * RemoveInProgressFiles removes all files that are no longer
 * in progress. This includes removing from the remote storage
 * and deleting the record from the IN_PROGRESS_FILES_TABLE.
 */
bool
RemoveInProgressFiles(char *location, bool isFull, bool isVerbose, List **deletedPaths)
{
	List	   *inProgressFileRecords = NIL;
	bool		hasRemainingFiles = GetInProgressFileRecords(location, isFull, &inProgressFileRecords);

	*deletedPaths = NIL;

	ListCell   *fileCell = NULL;

	foreach(fileCell, inProgressFileRecords)
	{
		InProgressFiles *entry = lfirst(fileCell);

		if (!TryLockTransactionOperationId(entry->operationId))
		{
			/* if we can't lock the operation id, it's still in progress */
			continue;
		}

		if (!FileInProgressRecordExists(entry->path))
		{
			/*
			 * Now that we have the lock, check if record exists. The
			 * operation could have completed successfully just after we have
			 * read the entry, in which case the record might be now gone.
			 */
			continue;
		}

		ereport(isVerbose ? INFO : LOG,
				(errmsg("deleting unused file %s%s",
						entry->path,
						entry->isPrefix ? "s" : "")));

		if (entry->isPrefix)
		{
			DeleteRemotePrefix(entry->path);
		}
		else
		{
			DeleteRemoteFile(entry->path);
		}

		DeleteInProgressFileRecord(entry->path);

		*deletedPaths = lappend(*deletedPaths, entry->path);
	}

	return hasRemainingFiles;
}


/*
 * GetInProgressFileRecords gets a list of paths that are eligible for
 * deletion from the IN_PROGRESS_FILES_TABLE.
 */
static bool
GetInProgressFileRecords(char *location, bool isFull, List **fileRecords)
{
	bool		hasRemainingFiles = false;

	*fileRecords = NIL;

	/* switch to schema owner, we assume callers checked permissions */
	Oid			savedUserId = InvalidOid;
	int			savedSecurityContext = 0;

	GetUserIdAndSecContext(&savedUserId, &savedSecurityContext);
	SetUserIdAndSecContext(ExtensionOwnerId(PgLakeEngine),
						   SECURITY_LOCAL_USERID_CHANGE);

	MemoryContext callerContext = CurrentMemoryContext;

	/* filter anything under "location" */
	StringInfo	locationFilter = makeStringInfo();

	appendStringInfo(locationFilter, "%s%%", location);

	StringInfo	query = makeStringInfo();

	appendStringInfo(query,
					 "    SELECT path, operation_id, is_prefix "
					 "    FROM " PG_LAKE_ENGINE_NSP "." IN_PROGRESS_FILES_TABLE " "
					 "    WHERE path ILIKE $1");

	DECLARE_SPI_ARGS(1);
	SPI_ARG_VALUE(1, TEXTOID, locationFilter->data, false);

	if (!isFull)
	{
		appendStringInfo(query,
						 "    LIMIT " PG_LAKE_TOSTRING(PER_LOOP_IN_PROGRESS_FILE_CLEANUP_LIMIT));
	}

	SPI_START();

	bool		readOnly = false;

	SPI_EXECUTE(query->data, readOnly);

	for (int rowIndex = 0; rowIndex < SPI_processed; rowIndex++)
	{
		bool		isNull;
		MemoryContext spiContext = MemoryContextSwitchTo(callerContext);

		InProgressFiles *entry = palloc0(sizeof(InProgressFiles));

		entry->path = GET_SPI_VALUE(TEXTOID, rowIndex, Anum_in_progress_path, &isNull);
		entry->operationId = GET_SPI_VALUE(INT8OID, rowIndex, Anum_in_progress_operation_id, &isNull);
		entry->isPrefix = GET_SPI_VALUE(BOOLOID, rowIndex, Anum_in_progress_is_prefix, &isNull);

		*fileRecords = lappend(*fileRecords, entry);
		MemoryContextSwitchTo(spiContext);
	}

	SPI_END();

	SetUserIdAndSecContext(savedUserId, savedSecurityContext);

	if (isFull)
	{
		hasRemainingFiles = false;
	}
	else
	{
		hasRemainingFiles = list_length(*fileRecords) >= PER_LOOP_IN_PROGRESS_FILE_CLEANUP_LIMIT;
	}

	return hasRemainingFiles;
}


/*
* DeleteInProgressFileRecord deletes the record from the
* IN_PROGRESS_FILES_TABLE.
*/
void
DeleteInProgressFileRecord(char *path)
{
	StringInfo	queryBuf = makeStringInfo();

	appendStringInfo(queryBuf,
					 "DELETE FROM " PG_LAKE_ENGINE_NSP "." IN_PROGRESS_FILES_TABLE " "
					 "WHERE path = $1");

	DECLARE_SPI_ARGS(1);
	SPI_ARG_VALUE(1, TEXTOID, path, false);

	/* switch to schema owner, we assume callers checked permissions */
	Oid			savedUserId = InvalidOid;
	int			savedSecurityContext = 0;

	GetUserIdAndSecContext(&savedUserId, &savedSecurityContext);
	SetUserIdAndSecContext(ExtensionOwnerId(PgLakeEngine),
						   SECURITY_LOCAL_USERID_CHANGE);

	SPI_START();
	SPIPlanPtr	queryPlan = GetCachedQueryPlan(queryBuf->data, spiArgCount, spiArgTypes);

	bool		readOnly = false;

	/* IN_PROGRESS_FILES_TABLE doesn't have any triggers */
	bool		fireTriggers = false;

	/* always read from the  latest snapshot */
	Snapshot	snapshot = GetLatestSnapshot();
	int			spiResult = SPI_execute_snapshot(queryPlan,
												 spiArgValues, spiArgNulls,
												 snapshot,
												 InvalidSnapshot,
												 readOnly, fireTriggers, 0);

	/* Check result */
	if (spiResult != SPI_OK_DELETE)
	{
		elog(ERROR, "SPI_execute_snapshot returned %s", SPI_result_code_string(spiResult));
	}

	SPI_END();

	SetUserIdAndSecContext(savedUserId, savedSecurityContext);

	CommandCounterIncrement();
}


/*
 * ReplaceInProgressPrefixPathWithFullPaths replaces an in-progress prefix path with the actual file paths.
 */
void
ReplaceInProgressPrefixPathWithFullPaths(char *prefixPath, List *fullPaths)
{
	/*
	 * Let's do a sanity check, ensuring fullPaths are prefixed with
	 * prefixPath
	 */

#ifdef USE_ASSERT_CHECKING
	ListCell   *pathCell = NULL;

	Assert(list_length(fullPaths) > 0);
	foreach(pathCell, fullPaths)
	{
		char	   *path = (char *) lfirst(pathCell);

		elog(DEBUG3, "expanding prefix %s with path: %s", prefixPath, path);
		Assert(strncmp(path, prefixPath, strlen(prefixPath)) == 0);
	}
#endif

	/*
	 * Let's replace the in-progress path with the actual file paths with a
	 * modifying CTE: WITH del_prefix (DELETE FROM in_progress_files WHERE
	 * path  = 'prefixPath') INSERT INTO in_progress_files (path,
	 * operation_id, is_prefix) SELECT path, operation_id, is_prefix FROM
	 * del_prefix;
	 */

	StringInfo	query = makeStringInfo();

	appendStringInfo(query,
					 "WITH del_prefix AS (DELETE FROM " PG_LAKE_ENGINE_NSP "." IN_PROGRESS_FILES_TABLE " "
					 "WHERE path = $1 AND is_prefix = TRUE RETURNING operation_id) "
					 "INSERT INTO " PG_LAKE_ENGINE_NSP "." IN_PROGRESS_FILES_TABLE " "
					 "(path, operation_id, is_prefix) "
					 "SELECT p, dp.operation_id, FALSE FROM del_prefix dp, unnest($2::text[]) p RETURNING *");

	/* switch to schema owner, we assume callers checked permissions */
	SPI_START_EXTENSION_OWNER(PgLakeEngine);

	DECLARE_SPI_ARGS(2);
	SPI_ARG_VALUE(1, TEXTOID, prefixPath, false);
	SPI_ARG_VALUE(2, TEXTARRAYOID, StringListToArray(fullPaths), false);

	SPIPlanPtr	queryPlan = GetCachedQueryPlan(query->data, spiArgCount, spiArgTypes);

	bool		readOnly = false;

	/* IN_PROGRESS_FILES_TABLE doesn't have any triggers */
	bool		fireTriggers = false;

	/* always read from the  latest snapshot */
	Snapshot	snapshot = GetLatestSnapshot();
	int			spiResult = SPI_execute_snapshot(queryPlan,
												 spiArgValues, spiArgNulls,
												 snapshot,
												 InvalidSnapshot,
												 readOnly, fireTriggers, 0);

	/* Check result */
	if (spiResult != SPI_OK_INSERT_RETURNING)
	{
		elog(ERROR, "SPI_execute_snapshot returned %s", SPI_result_code_string(spiResult));
	}

#ifdef USE_ASSERT_CHECKING

	/*
	 * must insert all paths if we are not in the same tx that created
	 * extension
	 */
	Assert(!InProgressTableVisibleToExternalTx() || SPI_processed == list_length(fullPaths));
#endif

	SPI_END();
}


/*
 * InsertInProgressFileRecord adds a path into the in progress file table
 * for internal bookkeeping.
 */
void
InsertInProgressFileRecord(char *path)
{
	InsertInProgressFileRecordExtended(path, false, false);
}


/*
 * InsertInProgressFileRecordExtended adds a path into the in progress file table
 * for internal bookkeeping with more options.
 *
 * "isPrefix": indicates whether the path is a prefix (directory) or a full path (file).
 * "deferDeletion": indicates whether the deletion of the record should be deferred until
 *                  the end of the transaction for iceberg tables.
 */
void
InsertInProgressFileRecordExtended(char *path, bool isPrefix, bool deferDeletion)
{
	if (XactReadOnly)
	{
		/* this check covers RecoveryInProgress() as well */
		return;
	}
	else if (!InProgressTableVisibleToExternalTx())
	{
		/*
		 * We are most likely in the same transaction that created the
		 * extension. The in-progress table is not visible to the external
		 * transactions until the transaction commits.
		 *
		 * Caching this value is not trivial, you should also take
		 * drop&recreate extension.
		 */
		return;
	}

	int64		operationId = GetTransactionOperationId();

	/*
	 * Prevent concurrent vacuum to operate on this file.
	 */
	LockTransactionOperationId(operationId);


	StringInfo	query = makeStringInfo();

	appendStringInfo(query,
					 "insert into " PG_LAKE_ENGINE_NSP "." IN_PROGRESS_FILES_TABLE " "
					 "(path, operation_id, is_prefix) "
					 "values (%s," INT64_FORMAT ", %s)",
					 quote_literal_cstr(path), operationId,
					 isPrefix ? "true" : "false");

	StringInfo	runAttachedQuery = makeStringInfo();

	appendStringInfo(runAttachedQuery,
					 "select * from extension_base.run_attached($$%s$$)",
					 query->data);


	/* switch to schema owner, we assume callers checked permissions */
	Oid			savedUserId = InvalidOid;
	int			savedSecurityContext = 0;

	GetUserIdAndSecContext(&savedUserId, &savedSecurityContext);
	SetUserIdAndSecContext(ExtensionOwnerId(PgLakeEngine),
						   SECURITY_LOCAL_USERID_CHANGE);

	SPI_START();

	bool		readOnly = false;

	SPI_execute(runAttachedQuery->data, readOnly, 0);

	if (SPI_processed != 1)
	{
		ereport(ERROR, (errmsg("failed to insert in progress file record")));
	}

	SPI_END();

	SetUserIdAndSecContext(savedUserId, savedSecurityContext);

	/*
	 * Huh? We just inserted the record, why are we removing it?
	 *
	 * This is central to the design of the in-progress files cleanup. We
	 * insert the record via an external transaction, so the record becomes
	 * immediately visible to the other transactions, including this backend.
	 *
	 * By removing the record from the in-progress table, we ensure that if
	 * this transaction commits, the record is removed. If the transaction
	 * aborts, the record (and the file) is removed by the vacuum.
	 */
	if (!deferDeletion)
		DeleteInProgressFileRecord(path);
}


/*
* InProgressTableVisibleToExternalTx checks if the in-progress
* table is visible to the external transactions.
*
* This is only useful when the in-progress table is created
* in the same transaction that is inserting the records.
*/
static bool
InProgressTableVisibleToExternalTx(void)
{
	StringInfo	query = makeStringInfo();

	appendStringInfo(query,
					 "SELECT * FROM extension_base.run_attached($$"
					 "SELECT 1 FROM information_schema.tables "
					 "WHERE table_schema OPERATOR(pg_catalog.=) '%s' "
					 "AND table_name OPERATOR(pg_catalog.=) '%s'$$);",
					 PG_LAKE_ENGINE_NSP, IN_PROGRESS_FILES_TABLE);

	bool		visible = false;

	/* switch to schema owner, we assume callers checked permissions */
	Oid			savedUserId = InvalidOid;
	int			savedSecurityContext = 0;

	GetUserIdAndSecContext(&savedUserId, &savedSecurityContext);
	SetUserIdAndSecContext(ExtensionOwnerId(PgLakeEngine),
						   SECURITY_LOCAL_USERID_CHANGE);

	SPI_START();

	bool		readOnly = true;

	SPI_execute(query->data, readOnly, 0);

	if (SPI_processed == 1)
	{

		bool		isNull;
		char	   *commandTag = GET_SPI_VALUE(TEXTOID, 0, 2, &isNull);

		/*
		 * run_attached returns the command tag of the query. We asked for
		 * whether the table exists, so the command tag should be "SELECT 1".
		 */
		if (strcmp(commandTag, "SELECT 1") == 0)
		{
			visible = true;
		}
	}

	SPI_END();

	SetUserIdAndSecContext(savedUserId, savedSecurityContext);

	return visible;
}

/*
* GetTransactionOperationId returns the operation id for the current
* transaction. If the operation id is not set, it gets a new operation
* id from the sequence.
*/
static int64
GetTransactionOperationId(void)
{
	if (TransactionOperationId != 0)
	{
		/* already set in this transaction */
		return TransactionOperationId;
	}

	if (!TransactionCallbackRegistered)
	{
		RegisterXactCallback(InProgressFileTransactionCallback, NULL);
		TransactionCallbackRegistered = true;
	}

	TransactionOperationId = GetNextTransactionOperationId();

	return TransactionOperationId;
}


/*
* GetNextTransactionOperationId returns the next operation id from the
* OPERATION_ID_SEQUENCE.
*/
static int64
GetNextTransactionOperationId(void)
{
	bool		missingOk = false;
	Oid			namespaceId = get_namespace_oid(PG_LAKE_ENGINE_NSP, missingOk);
	Oid			sequenceId = get_relname_relid(OPERATION_ID_SEQUENCE, namespaceId);

	return nextval_internal(sequenceId, false);
}



/*
* InProgressFileTransactionCallback is called when a transaction
* is committed or aborted. It resets the TransactionOperationId
* to 0.
*/
static void
InProgressFileTransactionCallback(XactEvent event, void *arg)
{
	if (event == XACT_EVENT_ABORT || event == XACT_EVENT_PARALLEL_ABORT ||
		event == XACT_EVENT_COMMIT || event == XACT_EVENT_PARALLEL_COMMIT
		|| event == XACT_EVENT_PREPARE)
	{
		/*
		 * We rely on the fact that the operationId is set uniquely for each
		 * transaction. Hence, we call it TransactionOperationId.
		 *
		 * Whenever a transaction needs a new operationId, it will use
		 * GetTransactionOperationId(). So, we'll get the next operationId
		 * from OPERATION_ID_SEQUENCE when needed.
		 */
		TransactionOperationId = 0;
	}
}

/*
 * LockTransactionOperationId takes an exclusive lock on the
 * operationId.
 */
static void
LockTransactionOperationId(int64 operationId)
{
	LOCKTAG		tag;
	const bool	sessionLock = false;
	const bool	dontWait = false;

	SET_LOCKTAG_IN_PROGRESS_OPERATION_ID(tag, operationId);
	(void) LockAcquire(&tag, ExclusiveLock, sessionLock, dontWait);
}

/*
 * TryLockTransactionOperationId takes an exclusive lock
 * on the given operationId. If the lock is not available,
 * it returns false.
 */
static bool
TryLockTransactionOperationId(int64 operationId)
{
	LOCKTAG		tag;
	const bool	sessionLock = false;
	const bool	dontWait = true;

	SET_LOCKTAG_IN_PROGRESS_OPERATION_ID(tag, operationId);
	LockAcquireResult lockResult =
		LockAcquire(&tag, ExclusiveLock, sessionLock, dontWait);

	return (lockResult != LOCKACQUIRE_NOT_AVAIL);
}


/*
* FileInProgressRecordExists checks if the given path
* exists in the IN_PROGRESS_FILES_TABLE.
*/
static bool
FileInProgressRecordExists(char *path)
{
	bool		exists = false;
	StringInfo	queryBuf = makeStringInfo();

	appendStringInfo(queryBuf,
					 "SELECT 1 FROM " PG_LAKE_ENGINE_NSP "." IN_PROGRESS_FILES_TABLE " "
					 "WHERE path = $1");

	DECLARE_SPI_ARGS(1);
	SPI_ARG_VALUE(1, TEXTOID, path, false);

	/* switch to schema owner, we assume callers checked permissions */
	Oid			savedUserId = InvalidOid;
	int			savedSecurityContext = 0;

	GetUserIdAndSecContext(&savedUserId, &savedSecurityContext);
	SetUserIdAndSecContext(ExtensionOwnerId(PgLakeEngine),
						   SECURITY_LOCAL_USERID_CHANGE);

	SPI_START();
	SPIPlanPtr	queryPlan = GetCachedQueryPlan(queryBuf->data, spiArgCount, spiArgTypes);

	bool		readOnly = true;

	/* IN_PROGRESS_FILES_TABLE doesn't have any triggers */
	bool		fireTriggers = false;

	/* always read from the  latest snapshot */
	Snapshot	snapshot = GetLatestSnapshot();

	int			spiResult =
		SPI_execute_snapshot(queryPlan,
							 spiArgValues, spiArgNulls,
							 snapshot,
							 InvalidSnapshot,
							 readOnly, fireTriggers, 0);

	/* Check result */
	if (spiResult != SPI_OK_SELECT)
	{
		elog(ERROR, "SPI_execute_snapshot returned %s", SPI_result_code_string(spiResult));
	}

	if (SPI_processed > 0)
	{
		bool		isNull = false;

		exists = GET_SPI_VALUE(BOOLOID, 0, 1, &isNull);
	}

	SPI_END();

	SetUserIdAndSecContext(savedUserId, savedSecurityContext);

	return exists;
}
